# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/12 16:58 
@Author  : ZhangShenao 
@File    : 1.使用PydanticOutputParser.py 
@Desc    : 使用PydanticOutputParser,将LLM的输出解析成结构化的数据,如json
"""
import os

import dotenv
import pandas as pd
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

from flower_description_output import FlowerDescriptionOutput

# 加载环境变量
dotenv.load_dotenv()

# 创建LLM
llm = ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE'))

# 创建空的DataFrame,用于存储结果
df = pd.DataFrame(columns=['flower_type', 'price', 'description', 'reason'])

# 创建PydanticOutputParser输出解析器,并指定pydantic结构
parser = PydanticOutputParser(pydantic_object=FlowerDescriptionOutput)

# 编排Prompt
template = '''你是一位专业的鲜花店文案撰写员。
对于价格为 {price} 元的 {flower_type}，请帮我生成一个吸引人的简短的中文文案。
{format_instruction}
'''
prompt = ChatPromptTemplate.from_template(template).partial(format_instruction=parser.get_format_instructions())

# 构造Chain
chain = prompt | llm | parser

# 构造数据
flower_types = ['玫瑰', '百合', '康乃馨']
prices = [50, 30, 20]

# 调用Chain
for flower_type, price in zip(flower_types, prices):
    # 获取LLM的输出结果
    result = chain.invoke({'flower_type': flower_type, 'price': price})

    # 将输出的Pydantic对象转换为字典
    result_dict = result.dict()

    # 将解析后的数据添加到DataFrame中
    df.loc[len(df)] = result_dict

# 打印完整结果
print(df.to_dict(orient='records'))
